#include "CustomBuffer.h"
#include "../../Core/CoreGlobals.h"
#include "../../Utils/LogUtil.h"
#include "../../Templates/NewObject.h"
#include "../RenderState.h"
#include "../../Game/Process/RenderProcess.h"

CustomBuffer::Builder::Builder() noexcept = default;
CustomBuffer::Builder::~Builder() noexcept = default;
CustomBuffer::Builder::Builder(CustomBuffer::Builder const& rhs) noexcept = default;
CustomBuffer::Builder::Builder(CustomBuffer::Builder && rhs) noexcept = default;
CustomBuffer::Builder& CustomBuffer::Builder::operator=(CustomBuffer::Builder const& rhs) noexcept = default;
CustomBuffer::Builder& CustomBuffer::Builder::operator=(CustomBuffer::Builder && rhs) noexcept = default;

RefCountPtr<CustomBuffer> CustomBuffer::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<CustomBuffer>(NewObject<CustomBuffer>(*mImpl));
}

RefCountPtr<CustomBuffer> CustomBuffer::Builder::Build(RenderProcess * process, const std::string & name)
{
	assert(IsInRenderThread());
	CustomBuffer* customBuffer = NewObject<CustomBuffer>(*mImpl);
	process->AddBuffer(name, customBuffer);
	return RefCountPtr<CustomBuffer>(customBuffer);
}

void CustomBuffer::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemoveBuffer(mName, this);
		}
		DeleteObject<CustomBuffer>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemoveBuffer(mName, this);
			}
			DeleteObject<CustomBuffer>(this);
			});
	}
}

CustomBuffer::Builder& CustomBuffer::Builder::SetMultiple(bool isMultiple) noexcept
{
	mImpl->isMultiple = isMultiple;
	return *this;
}

CustomBuffer::Builder& CustomBuffer::Builder::SetFlags(VkBufferCreateFlags flags) noexcept
{
	mImpl->flags = flags;
	return *this;
}

CustomBuffer::Builder& CustomBuffer::Builder::SetSize(VkDeviceSize size) noexcept
{
	mImpl->size = size;
	return *this;
}

CustomBuffer::Builder& CustomBuffer::Builder::SetUsage(VkBufferUsageFlags usage) noexcept
{
	mImpl->usage = usage;
	return *this;
}

CustomBuffer::Builder& CustomBuffer::Builder::SetSharingMode(VkSharingMode sharingMode) noexcept
{
	mImpl->sharingMode = sharingMode;
	return *this;
}

CustomBuffer::Builder& CustomBuffer::Builder::AddQueueFamily(VkQueueFlagBits queueFlag) noexcept
{
	mImpl->queueFamilys.push_back(queueFlag);
	return *this;
}

CustomBuffer::Builder& CustomBuffer::Builder::SetMemoryProperty(VkMemoryPropertyFlags memoryProperty) noexcept
{
	mImpl->memoryProperty = memoryProperty;
	return *this;
}

CustomBuffer::Builder& CustomBuffer::Builder::SetData(void* data) noexcept
{
	mImpl->data = data;
	return *this;
}

CustomBuffer::CustomBuffer(BuilderDetails & builderDetails)
{
	mIsMultiple = builderDetails.isMultiple;

	mFlags = builderDetails.flags;
	mSize = builderDetails.size;
	mUsage = builderDetails.usage;
	mSharingMode = builderDetails.sharingMode;
	mQueueFamilys.assign(builderDetails.queueFamilys.begin(), builderDetails.queueFamilys.end());
	if (mQueueFamilys.size() > 1)
	{
		// multi queue
		assert(mSharingMode == VK_SHARING_MODE_CONCURRENT);
	}
	else
	{
		assert(mSharingMode == VK_SHARING_MODE_EXCLUSIVE);
	}
	mMemoryProperty = builderDetails.memoryProperty;

	mBuffers.resize(GetFrameCount(), VK_NULL_HANDLE);
	mBufferMemorys.resize(GetFrameCount(), VK_NULL_HANDLE);

	for (uint32 i = 0; i < GetFrameCount(); ++i)
	{
		GVulkanRHI->CreateBuffer(mBuffers[i], mBufferMemorys[i], mSize, mUsage, mMemoryProperty, mSharingMode, mQueueFamilys);
	}


	// Decide use which way to write data
	if (builderDetails.data != nullptr)
	{
		// Only propertyFlags is VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT we can vkMapMemory
		// If propertyFlags is VK_MEMORY_PROPERTY_HOST_COHERENT_BIT we need to use vkFlushMappedMemoryRanges and vkInvalidateMappedMemoryRanges to write data
		if (mMemoryProperty & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT)
		{
			if (mMemoryProperty & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
			{
				for (uint32 i = 0; i < GetFrameCount(); ++i)
				{
					void* data;
					VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), mBufferMemorys[i], 0, mSize, 0, &data));
					memcpy(data, builderDetails.data, mSize);
					vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), mBufferMemorys[i]);
				}
			}
			else
			{
				for (uint32 i = 0; i < GetFrameCount(); ++i)
				{
					void* data;
					VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), mBufferMemorys[i], 0, mSize, 0, &data));
					memcpy(data, builderDetails.data, mSize);

					// Flush to make changes visible to the host
					VkMappedMemoryRange mappedMemoryRange;
					mappedMemoryRange.sType = VK_STRUCTURE_TYPE_MAPPED_MEMORY_RANGE;
					mappedMemoryRange.memory = mBufferMemorys[i];
					mappedMemoryRange.offset = 0;
					mappedMemoryRange.size = mSize;
					vkFlushMappedMemoryRanges(GVulkanRHI->GetLogicalDevice(), 1, &mappedMemoryRange);
					vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), mBufferMemorys[i]);
				}
			}
		}
		else if (mMemoryProperty & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)
		{
			// We can only use stage buffer to copy data on device
			// Usage need to set VK_BUFFER_USAGE_TRANSFER_DST_BIT
			assert(mUsage & VK_BUFFER_USAGE_TRANSFER_DST_BIT);

			VkBuffer stagingBuffer;
			VkDeviceMemory stagingMemory;

			GVulkanRHI->CreateBuffer(stagingBuffer, stagingMemory, mSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

			void* data;
			VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, 0, mSize, 0, &data));
			memcpy(data, builderDetails.data, mSize);
			vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory);

			VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();


			for (uint32 i = 0; i < GetFrameCount(); ++i)
			{
				VkBufferCopy bufferCopy{};
				bufferCopy.srcOffset = 0;
				bufferCopy.dstOffset = 0;
				bufferCopy.size = mSize;
				vkCmdCopyBuffer(commandBuffer, stagingBuffer, mBuffers[i], 1, &bufferCopy);
			}


			GRenderState->FlushSingleCommand(commandBuffer);

			vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), stagingBuffer, nullptr);
			vkFreeMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, nullptr);
		}
		else
		{
			//TODO: more ways to copy data
			assert(0);
		}
	}
}